"""Functions for the flux computations"""

from hogs.solvers.flux.flux_functions import FluxFunction, WAFFlux, LxfFlux

from hogs.solvers.riemann.euler import euler
from hogs.solvers.riemann.euler import reuler
from hogs.solvers.riemann.euler import hllc
from hogs.solvers.riemann import solver_types

from hogs.solvers.limiters import limiters, waf_limiters

import numpy
from optparse import OptionGroup

class RiemannSolverOptions(object):
    """Options for the Riemann solvers"""    
    def __init__(self, gamma=1.4, rsolver_type=0):
        self.defaults = dict(rsolver_type=rsolver_type,
                             gamma=gamma)

    def get_options(self, parser):
        """Set options for the RiemannSolverFlux"""
        opt = OptionGroup(parser, "Options for the Riemann solver")

        opt.add_option("--gamma", action="store", type="float",
                       dest="gamma", default=self.defaults["gamma"],
                       help="Adiabatic constant gamma")

        opt.add_option("--rsolver", action="store", type="int",
                       dest="rsolver_type",
                       default=self.defaults.get("rsolver_type"),
                       help="The Riemann solver to use:\t\t\t\t" + \
                       "".join(['%d - %-51s'%(d,s[0]) for d,s in
                                enumerate(solver_types.euler_solvers)]))

        hllc_opt = OptionGroup(parser, """Optional options for the HLLC Riemann solver""")

        hllc_opt.add_option("--wave_speed", action="store", type="int",
                       dest="wave_speed", default=0,
                       help="Wave speed estimator for the HLLC solver:\t\t" +\
                        "".join(['%d - %-51s'%(d,s[0]) for d,s in
                                 enumerate(hllc.euler_wave_speeds)]))

        return opt, hllc_opt

    def set_options(self, fluxfunc, options):
        """Set the Riemann solver"""
        options = options or self.defaults

        fluxfunc.gamma = gamma = options.get("gamma")

        # get the Riemann solver type 
        rsolver_type = options.get("rsolver_type")
        fluxfunc.rsolver = solver_types.euler_solvers[rsolver_type][1](gamma)

        # set the wave spped estimator for the HLLC solver
        if solver_types.euler_solvers[rsolver_type][0] == "HLLC":
            wave_speed = options.get("wave_speed")
            estimator = hllc.euler_wave_speeds[wave_speed][1](gamma)

            fluxfunc.rsolver.set_speed_estimator( estimator )    

###########################################################################
#`GodunovFluxEulerian`
###########################################################################
class GodunovFluxEulerian(FluxFunction):
    """Godunov's first order flux for the Euler equations.

    The Riemann solver used for the flux evaluation is defined in
    reuler.py and is basically a pythonised version of the code in
    Toro's book.

    """
    def __init__(self, gamma=1.4, rsolver_type=0):
        self.rsolver_opts = RiemannSolverOptions(gamma, rsolver_type)

    def get_options(self, parser):
        rsolver_opts = self.rsolver_opts.get_options(parser)
        return rsolver_opts

    def set_options(self, options):
        self.rsolver_opts.set_options(self, options)
        
    def compute_flux(self, grid, k):

        faces = grid.faces
        nb = grid.nb

        # compute the flux for each face
        for face in faces[nb:-nb]:

            # get the primitive variables for each adjoining cell
            wl, wr = face.get_primitive_variables()

            # store the flux for this face
            self.rsolver.set_eulerian_godunov_flux(wl, wr, face)

############################################################################w
# `WAFFlux` class
############################################################################w
class WAFFluxEulerian(WAFFlux):
    """FLux function for the WAF FLux (Toro).

    The Weighted average flux (WAF) can be thought of an extension to
    Godunov's first order flux. The general form of the flux used is
    given as:

    .. math::

    F{i + \frac{1}{2}} = \frac{1}{\Delta x}\int_{-\Delta x}^{\Delta x}
    F\left( U_{i+\frac{1}{2}}\left(x, \frac{\Delta t}{2} \right )
    \right )


    After performing the above integral on the structure of the Euler
    equations, the resulting flux form is:

    ..math ::

    F{i + \frac{1}{2}} = \frac{1}{2}(F_i + F_{i+1}) -
    \frac{1}{2}\sum_{k=1}^{N} sign(c_k) \Delta F_{i + \frac{1}{2}}^k

    """
    def __init__(self, gamma=1.4, rsolver_type=0, limiter_type=0):
        self.rsolver_opts = RiemannSolverOptions(gamma, rsolver_type)
        WAFFlux.__init__(self, limiter_type)

    def get_options(self, parser):
        rsolver_opts = self.rsolver_opts.get_options(parser)
        waf_opts = WAFFlux.get_options(self, parser)
        
        return rsolver_opts, waf_opts

    def set_options(self, options):
        """Set options for the WAF flux"""
        self.rsolver_opts.set_options(self, options)
        WAFFlux.set_options(self, options)

    def _get_data_state_vector(self):
        return numpy.zeros( shape=(3, 4) )

    def _get_wave_speed_vector(self):
        return numpy.zeros(3)

    def _get_jump(self, ql, qr):
        return qr[0] - ql[0]

    def _compute_unstable_flux(self, ql, qr):
        return 0.5 * (euler.ideal_gas_flux_q(self.gamma, ql) + \
                      euler.ideal_gas_flux_q(self.gamma, qr) )

    def _compute_delF(self, U, i):
        Fi = euler.ideal_gas_flux_q( self.gamma, U[:, i] )
        Fip1 = euler.ideal_gas_flux_q( self.gamma, U[:, i + 1] )

        return Fip1 - Fi

    def _compute_waf_states(self, ql, qr, face, U, S):
        wl = euler.ideal_gas_q2w(self.gamma, ql)
        wr = euler.ideal_gas_q2w(self.gamma, qr)

        #  the left and right states
        U[:, 0] = ql; U[:, 3] = qr

        # compute the states between the waves
        self.rsolver.compute_waf_states(wl, wr, face, U, S)

class EulerLxf(LxfFlux):

    def _compute_unstable_flux(self, ql, qr):
        return 0.5 * ( euler.ideal_gas_flux_q(1.4, ql) + \
                       euler.ideal_gas_flux_q(1.4, qr) )

    def _compute_max_sound_speed(self):
        return 1.0
        
